#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import random
import typing

from mitmdump import DumpMaster, Options
from mitmproxy import http
from mitmproxy.http import HTTPFlow

# 在header中传递自定义代理ip的key值
KEY_PROXY_IP: str = 'proxy_ip'


def read_file(_path) -> str:
    """读取文件"""
    if os.path.exists(_path) is False:
        return None
    with open(_path, "r", encoding='utf-8') as f:  # 打开文件
        data = f.read()
        f.close()
        return data


class Proxy(object):
    # 当前请求使用的代理ip
    curr_proxy_ip: str = None
    # 默认的代理ip列表
    proxy_list: list = None

    def __init__(self):
        _txt: str = read_file('custom_proxy.txt')
        if _txt:
            self.proxy_list = _txt.split('\n')
        else:
            self.proxy_list = ['']
        pass

    def http_connect(self, flow: http.HTTPFlow):
        pass

    def request(self, flow: http.HTTPFlow) -> None:
        if flow.request.method == "CONNECT":
            return

        user_agent: str = flow.request.headers.get('User-Agent')
        proxy_ip: str = flow.request.headers.get(KEY_PROXY_IP)
        print(f'{KEY_PROXY_IP}: {proxy_ip} user_agent: {user_agent}\n')

        # 正式的二次代理请求中不要将代理信息传过去
        if proxy_ip:
            flow.request.headers.pop(KEY_PROXY_IP)

        proxy_info = self.get_proxy_tuple(proxy_ip)
        if proxy_info and flow.live:
            print(f'proxy_info: {proxy_info}')
            flow.live.change_upstream_proxy_server((proxy_info[0], proxy_info[1]))
            self.curr_proxy_ip = proxy_info[2]

    def response(self, flow: HTTPFlow):
        # 将本次请求中的代理信息回传，如果多个请求需要使用同一个代理，则下次请求可将该自定义代理传过来
        flow.response.headers[KEY_PROXY_IP] = self.curr_proxy_ip
        print(f'{flow.response.status_code}\n{flow.response.text}')
        pass

    def get_proxy_tuple(self, proxy_ip: str) -> typing.Tuple[str, str, str]:
        """
        获取代理信息
        :param proxy_ip: 如果有值则优先使用该代理，否则获取一个随机代理
        :return:
        """
        proxy_ip: str = proxy_ip or random.choice(self.proxy_list)
        if '//' in proxy_ip and ':' in proxy_ip:
            proxy_tuple = proxy_ip.split('//')[1].split(':')
            return proxy_tuple[0], proxy_tuple[1], proxy_ip
        else:
            print(f'proxy_ip 为空：{proxy_ip}')
            return None


addons = [
    Proxy()
]

if __name__ == '__main__':
    opts = Options(
        listen_host='0.0.0.0',
        listen_port=8080,
        scripts=None,
        dumper_filter='~m GET',
        mode='upstream:http://192.168.2.136:8222',  # 预设一个可访问的代理，后续程序中会动态替换
        # mode='upstream:SPEC',
        flow_detail=1,
        termlog_verbosity='info',
        show_clientconnect_log=False,
    )
    m = DumpMaster(opts)
    m.addons.add(*addons)
    m.run()
